#!/usr/bin/env python3
"""
Population Individual PBPK Modeling

This script generates realistic PK curves for individual subjects using
population-derived parameters, following the four-mode pattern from
individual_pbpk_modeling.py.

Model combinations:
- first_order + minimal_pbpk
- first_order + full_pbpk  
- adam + minimal_pbpk
- adam + full_pbpk

Author: BlackCat@CPPO
Version: 1.0.0
"""

import numpy as np
import matplotlib.pyplot as plt
import sys
import os
import json
from datetime import datetime

# Add src path to system path
project_root = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, project_root)

from demo.generate_pbpk_parameters import PopulationParameterGenerator
from demo.individual_pbpk_modeling import IndividualPBPKSimulator


class PopulationIndividualPBPK:
    """Individual PBPK modeling with population parameters"""
    
    def __init__(self, output_dir='../output/population_individual'):
        self.output_dir = output_dir
        self.parameter_generator = PopulationParameterGenerator()
        self.simulator = IndividualPBPKSimulator()
        self.ensure_output_dir()
    
    def ensure_output_dir(self):
        """Create output directory if it doesn't exist"""
        os.makedirs(self.output_dir, exist_ok=True)
    
    def run_all_models(self, subject_ids=None):
        """
        Run all four model combinations for specified subjects
        
        Args:
            subject_ids: List of subject IDs to model (None for all)
        """
        if subject_ids is None:
            subject_ids = [1, 5, 10, 25, 50, 75, 100]
        
        print("=" * 80)
        print("Population Individual PBPK Modeling - All Four Modes")
        print("=" * 80)
        
        # Define all four model combinations
        model_combinations = [
            ('first_order', 'minimal_pbpk'),
            ('first_order', 'full_pbpk'),
            ('adam', 'minimal_pbpk'),
            ('adam', 'full_pbpk')
        ]
        
        all_results = {}
        
        for subject_id in subject_ids:
            try:
                print(f"\nProcessing Subject {subject_id}...")
                
                # Get population parameters for this subject
                params = self.parameter_generator.get_individual_parameters(subject_id=subject_id)
                
                print(f"  Profile: {params['body_weight']:.1f}kg, {params['height']:.1f}cm, {params['age']:.1f}y, {params['sex']}")
                
                subject_results = {}
                
                for abs_model, dist_model in model_combinations:
                    try:
                        print(f"  Running {abs_model} + {dist_model}...")
                        
                        # Prepare ADAM parameters if needed
                        if abs_model == 'adam':
                            params['use_adam'] = True
                        
                        # Run simulation
                        result = self.simulator.run_complete_pbpk(abs_model, dist_model, params)
                        
                        # Calculate PK metrics
                        metrics = self.simulator.calculate_pk_metrics(
                            result['time_points'],
                            result['elimination']['plasma_concentration']
                        )
                        
                        # Store results
                        key = f"{abs_model}_{dist_model}"
                        subject_results[key] = {
                            'parameters': params,
                            'result': result,
                            'pk_metrics': metrics
                        }
                        
                        print(f"    Cmax: {metrics['C_max']:.3f} mg/L, Tmax: {metrics['T_max']:.1f}h, AUC: {metrics['AUC']:.1f} mg·h/L")
                        
                    except Exception as e:
                        print(f"    Error: {str(e)}")
                        continue
                
                all_results[f"subject_{subject_id}"] = subject_results
                
                # Save individual subject results
                self.save_subject_results(subject_id, subject_results)
                
            except Exception as e:
                print(f"  Failed to process subject {subject_id}: {str(e)}")
                continue
        
        # Generate plots
        self.plot_all_subjects(all_results)
        self.plot_model_comparisons(all_results)
        
        return all_results
    
    def save_subject_results(self, subject_id, results):
        """Save results for a single subject"""
        filename = os.path.join(self.output_dir, f'subject_{subject_id}_results.json')
        
        # Convert numpy arrays to lists for JSON serialization
        serializable_results = {}
        for model_key, data in results.items():
                    serializable_results[model_key] = {
                        'pk_metrics': data['pk_metrics'],
                        'time_points': data['result']['time_points'].tolist(),
                        'concentration': data['result']['elimination']['plasma_concentration'].tolist(),
                        'parameters': {k: float(v) if isinstance(v, (np.integer, np.floating)) else v 
                                     for k, v in data['parameters'].items()
                                     if not isinstance(v, (np.ndarray, dict))}
                    }
        
        with open(filename, 'w') as f:
            json.dump(serializable_results, f, indent=2)
    
    def plot_all_subjects(self, all_results):
        """Plot all subjects for each model combination"""
        model_combinations = [
            ('first_order', 'minimal_pbpk'),
            ('first_order', 'full_pbpk'),
            ('adam', 'minimal_pbpk'),
            ('adam', 'full_pbpk')
        ]
        
        fig, axes = plt.subplots(2, 2, figsize=(16, 12))
        axes = axes.flatten()
        
        colors = plt.cm.Set1(np.linspace(0, 1, len(all_results)))
        
        for idx, (abs_model, dist_model) in enumerate(model_combinations):
            ax = axes[idx]
            model_key = f"{abs_model}_{dist_model}"
            
            for i, (subject_key, subject_results) in enumerate(all_results.items()):
                if model_key in subject_results:
                    data = subject_results[model_key]
                    ax.plot(data['result']['time_points'], 
                           data['result']['elimination']['plasma_concentration'],
                           color=colors[i % len(colors)],
                           alpha=0.7, linewidth=2,
                           label=f"Subject {subject_key.split('_')[1]}")
            
            ax.set_xlabel('Time (h)')
            ax.set_ylabel('Plasma Concentration (mg/L)')
            ax.set_title(f'{abs_model.replace("_", " ").title()} + {dist_model.replace("_", " ").title()}')
            ax.grid(True, alpha=0.3)
            ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
        
        plt.tight_layout()
        filename = os.path.join(self.output_dir, 'all_subjects_all_models.png')
        plt.savefig(filename, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"Saved all subjects plot: {filename}")
    
    def plot_model_comparisons(self, all_results):
        """Plot comparison of models for each subject"""
        n_subjects = min(4, len(all_results))
        subject_keys = list(all_results.keys())[:n_subjects]
        
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        axes = axes.flatten()
        
        model_colors = {
            'first_order_minimal_pbpk': 'blue',
            'first_order_full_pbpk': 'red',
            'adam_minimal_pbpk': 'green',
            'adam_full_pbpk': 'purple'
        }
        
        for idx, subject_key in enumerate(subject_keys):
            ax = axes[idx]
            subject_results = all_results[subject_key]
            
            for model_key, color in model_colors.items():
                if model_key in subject_results:
                    data = subject_results[model_key]
                    ax.plot(data['result']['time_points'],
                           data['result']['elimination']['plasma_concentration'],
                           color=color, linewidth=2,
                           label=model_key.replace('_', ' + '))
            
            ax.set_xlabel('Time (h)')
            ax.set_ylabel('Plasma Concentration (mg/L)')
            ax.set_title(f'{subject_key.replace("_", " ").title()}')
            ax.grid(True, alpha=0.3)
            ax.legend()
        
        plt.tight_layout()
        filename = os.path.join(self.output_dir, 'model_comparisons_per_subject.png')
        plt.savefig(filename, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"Saved model comparisons plot: {filename}")
    
    def generate_summary_report(self, all_results):
        """Generate summary report"""
        report_file = os.path.join(self.output_dir, 'individual_summary_report.txt')
        
        with open(report_file, 'w') as f:
            f.write("Population Individual PBPK Modeling Report\n")
            f.write("=" * 50 + "\n")
            f.write(f"Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write(f"Subjects analyzed: {len(all_results)}\n")
            f.write(f"Output directory: {self.output_dir}\n\n")
            
            model_combinations = [
                ('first_order', 'minimal_pbpk'),
                ('first_order', 'full_pbpk'),
                ('adam', 'minimal_pbpk'),
                ('adam', 'full_pbpk')
            ]
            
            for abs_model, dist_model in model_combinations:
                model_key = f"{abs_model}_{dist_model}"
                f.write(f"\n{model_key.replace('_', ' + ').title()}:\n")
                f.write("-" * 40 + "\n")
                
                Cmax_values = []
                AUC_values = []
                Tmax_values = []
                
                for subject_key, subject_results in all_results.items():
                    if model_key in subject_results:
                        metrics = subject_results[model_key]['pk_metrics']
                        Cmax_values.append(metrics['C_max'])
                        AUC_values.append(metrics['AUC'])
                        Tmax_values.append(metrics['T_max'])
                
                if Cmax_values:
                    f.write(f"Cmax (mg/L): {np.mean(Cmax_values):.3f} ± {np.std(Cmax_values):.3f}\n")
                    f.write(f"AUC (mg·h/L): {np.mean(AUC_values):.1f} ± {np.std(AUC_values):.1f}\n")
                    f.write(f"Tmax (h): {np.mean(Tmax_values):.1f} ± {np.std(Tmax_values):.1f}\n")
                    f.write(f"Subjects: {len(Cmax_values)}\n")
        
        print(f"Summary report saved: {report_file}")


def main():
    """Main function"""
    print("Population Individual PBPK Modeling")
    print("=" * 50)
    
    # Create runner
    runner = PopulationIndividualPBPK()
    
    # Run all models for representative subjects
    results = runner.run_all_models(subject_ids=[1, 5, 10, 25, 50])
    
    # Generate summary report
    runner.generate_summary_report(results)
    
    print("\n" + "=" * 50)
    print("Modeling completed!")
    print(f"Results saved in: {runner.output_dir}")
    print("=" * 50)


if __name__ == "__main__":
    main()